// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/vector_map.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/vector_stream.h>

namespace kumo::pollux {
    VectorMap::VectorMap(BaseVector &alphabet)
        : alphabet_(&alphabet),
          isString_(
              alphabet_->type_kind() == TypeKind::VARCHAR ||
              alphabet_->type_kind() == TypeKind::VARBINARY),
          fixedWidth_(
              alphabet_->type()->isFixedWidth()
                  ? alphabet_->type()->cppSizeInBytes()
                  : kVariableWidth) {
        auto size = alphabet_->size();
        // We reserve the size. The assumption is that we run this on an alphabet of a
        // dictionary vector in preparation for adding elements, so the values are
        // expected to be distinct.
        if (isString_) {
            distinctStrings_.reserve(size);
        } else {
            distinctSet_.reserve(size);
        }
        for (auto i = 0; i < size; ++i) {
            addOne(*alphabet_, i, false);
        }
    }

    VectorMap::VectorMap(
        const TypePtr &type,
        memory::MemoryPool *pool,
        int32_t reserve)
        : isString_(
              type->kind() == TypeKind::VARCHAR ||
              type->kind() == TypeKind::VARBINARY),
          fixedWidth_(
              type->isFixedWidth() ? type->cppSizeInBytes() : kVariableWidth) {
        alphabetOwned_ = BaseVector::create(type, 0, pool);
        alphabet_ = alphabetOwned_.get();
        if (isString_) {
            distinctStrings_.reserve(reserve);
        } else {
            distinctSet_.reserve(reserve);
        }
    }

    vector_size_t VectorMap::addOne(
        const BaseVector &topVector,
        vector_size_t topIndex,
        bool insertToAlphabet) {
        const BaseVector *vector = &topVector;
        vector_size_t index = topIndex;
        if (topVector.encoding() != VectorEncoding::Simple::FLAT) {
            vector = topVector.wrapped_vector();
            index = topVector.wrapped_index(topIndex);
        }
        if (LIKELY(isString_)) {
            StringView string;
            bool isNull = vector->is_null_at(index);
            if (UNLIKELY(isNull)) {
                if (nullIndex_ != kNoNullIndex) {
                    return nullIndex_;
                }
            } else {
                if (LIKELY(vector->encoding() == VectorEncoding::Simple::FLAT)) {
                    string = vector->as_unchecked<FlatVector<StringView> >()->value_at(index);
                } else {
                    string = vector->as_unchecked<ConstantVector<StringView> >()->value_at(0);
                }
                auto it = distinctStrings_.find(string);
                if (it != distinctStrings_.end()) {
                    return it->second;
                }
            }
        } else {
            auto it = distinctSet_.find(VectorValueSetEntry{vector, index});
            if (it != distinctSet_.end()) {
                return it->index;
            }
        }
        int32_t newIndex;
        if (insertToAlphabet) {
            POLLUX_CHECK(alphabet_ == alphabetOwned_.get());
            newIndex = alphabet_->size();
            alphabet_->resize(newIndex + 1);
            alphabetSizes_.resize(newIndex + 1);
            alphabet_->copy(vector, newIndex, index, 1);
        } else {
            newIndex = topIndex;
            if (alphabetSizes_.size() < newIndex + 1) {
                alphabetSizes_.resize(newIndex + 1);
            }
        }
        const bool isNull = vector->is_null_at(index);
        if (isNull) {
            alphabetSizes_[newIndex] = 0;
            nullIndex_ = newIndex;
        } else if (isString_) {
            // Add 4 for length.
            alphabetSizes_[newIndex] = alphabet_->as_unchecked<FlatVector<StringView> >()
                                       ->value_at(newIndex)
                                       .size() +
                                       4;
        } else if (fixedWidth_ == kVariableWidth) {
            Scratch scratch;
            ScratchPtr<vector_size_t, 1> indicesHolder(scratch);
            ScratchPtr<vector_size_t *, 1> sizesHolder(scratch);
            auto sizeIndices = indicesHolder.get(1);
            sizeIndices[0] = newIndex;
            auto sizes = sizesHolder.get(1);
            alphabetSizes_[newIndex] = 0;
            sizes[0] = &alphabetSizes_[newIndex];
            getVectorSerde()->estimateSerializedSize(
                alphabet_,
                melon::Range<const vector_size_t *>(sizeIndices, 1),
                sizes,
                scratch);
        }
        if (isString_) {
            if (!isNull) {
                distinctStrings_[alphabet_->as_unchecked<FlatVector<StringView> >()
                    ->value_at(newIndex)] = newIndex;
            }
        } else {
            distinctSet_.insert(VectorValueSetEntry{alphabet_, newIndex});
        }
        return newIndex;
    }

    void VectorMap::addMultiple(
        BaseVector &vector,
        melon::Range<const vector_size_t *> rows,
        bool insertToAlphabet,
        vector_size_t *ids) {
        auto size = rows.size();
        for (auto i = 0; i < size; ++i) {
            ids[i] = addOne(vector, rows[i], insertToAlphabet);
        }
    }
} // namespace kumo::pollux
